#include "math.h"
#include "DependencyEncoder.h"
#include "ObjectReader.h"
#include "DependencyInstance.h"
#include <ctime>

DependencyEncoder::DependencyEncoder(int maxLength, DependencyPipe *pipe, char *trainforest, Parameters *param){
	in = new ObjectReader(trainforest);
	this->pipe = pipe;
	this->param = param;
	probs = vector<vector<vector<double> > >(maxLength);
	for(int i = 0; i < maxLength; i++){
		probs[i] = vector<vector<double> >(maxLength);
		for(int j = 0; j < maxLength; j++){
			probs[i][j] = vector<double>(2);
		}
	}

	if(pipe->labeled){
		int len = (int)(pipe->types.size());
		nt_probs = vector<vector<vector<vector<double> > > >(maxLength);
		for(int i = 0; i < maxLength; i++){
			nt_probs[i] = vector<vector<vector<double> > >(len);
			for(int j = 0; j < len; j++){
				nt_probs[i][j] = vector<vector<double> >(2); 
				for(int k = 0; k < 2; k++){
					nt_probs[i][j][k] = vector<double>(2);
				}
			}
		}
	}
}

DependencyEncoder::~DependencyEncoder(){
	delete(in);
}

int DependencyEncoder::getKey(int s, int t, int dir, int comp, int length){
	int key = s;
	key = key * length + t;
	key = key * 2 + dir;
	key = key * 2 + comp;
	return key;
}

double DependencyEncoder::getZ(int length, double *beta){
	int key;

	for(int i = 0; i < length; i++){
		key = getKey(i, i, 0, 1, length);
		beta[key] = 0.0;
		key = getKey(i, i, 1, 1, length);
		beta[key] = 0.0;
	}
	
	for(int j = 1; j < length; j++){
		for(int s = 0; s + j < length; s++){
			int t = s + j;
			double prodProb_st = probs[s][t][0];
			double prodProb_ts = probs[s][t][1];
			
			//init beta
			//incomplete spans
			int key_st_0 = getKey(s, t, 0, 0, length);
			beta[key_st_0] = 0.0;
			int key_ts_0 = getKey(s, t, 1, 0, length);
			beta[key_ts_0] = 0.0;
			
			//complete spans
			int key_st_1 = getKey(s, t, 0, 1, length);
			beta[key_st_1] = 0.0;
			int key_ts_1 = getKey(s, t, 1, 1, length);
			beta[key_ts_1] = 0.0;

			bool flg_st_0 = true, flg_ts_0 = true;
			bool flg_st_1 = true, flg_ts_1 = true;

			for(int r = s; r < t; r++){
				// first is direction, second is complete
				// _s means s is the parent
				int key1 = getKey(s, r, 0, 1, length);
				int key2 = getKey(r + 1, t, 1, 1, length);
				
				beta[key_st_0] = logsumexp(beta[key_st_0], beta[key1] + beta[key2] + prodProb_st, flg_st_0);
				flg_st_0 = false;
				
				beta[key_ts_0] = logsumexp(beta[key_ts_0], beta[key1] + beta[key2] + prodProb_ts, flg_ts_0);
				flg_ts_0 = false;
			}
			
			for(int r = s; r <= t; r++){
				if(r != s){
					int key1 = getKey(s, r, 0, 0, length);
					int key2 = getKey(r, t, 0, 1, length);
					beta[key_st_1] = logsumexp(beta[key_st_1], beta[key1] + beta[key2], flg_st_1);
					flg_st_1 = false;
				}
				if(r != t){
					int key1 = getKey(s, r, 1, 1, length);
					int key2 = getKey(r, t, 1, 0, length);
					beta[key_ts_1] = logsumexp(beta[key_ts_1], beta[key1] + beta[key2], flg_ts_1);
					flg_ts_1 = false;
				}
			}
		}
	}
	
	int key1 = getKey(0, length - 1, 0, 1, length);
	int key2 = getKey(0, length - 1, 1, 1, length);
	return logsumexp(beta[key1], beta[key2], false);
}

void DependencyEncoder::inside_outside(int length, double *beta, double *alpha){
	int key;
	int end = length - 1;
	for(int d = 0; d < 2; d++){
		for(int c = 0 ; c < 2; c++){
			key = getKey(0, end, d, c, length);
			alpha[key] = 0.0;
		}
	}

	for(int j = end; j >= 1; j--){
		for(int s = 0; s + j < length; s++){
			int t = s + j;

			//init alpha
			//incomplete spans
			int key_st_0 = getKey(s, t, 0, 0, length);
			alpha[key_st_0] = 0.0;
			int key_ts_0 = getKey(s, t, 1, 0, length);
			alpha[key_ts_0] = 0.0;
			
			//complete spans
			int key_st_1 = getKey(s, t, 0, 1, length);
			alpha[key_st_1] = 0.0;
			int key_ts_1 = getKey(s, t, 1, 1, length);
			alpha[key_ts_1] = 0.0;

			bool flg_st_0 = true, flg_ts_0 = true;
			bool flg_st_1 = true, flg_ts_1 = true;

			for(int r = 0; r < s; r++){
				double prodProb_rt = probs[r][t][0];
				double prodProb_tr = probs[r][t][1];

				//alpha[s][t][0][1]
				int key_b = getKey(r, s, 0, 0, length);
				int key_a = getKey(r, t, 0, 1, length);
				alpha[key_st_1] = logsumexp(alpha[key_st_1], beta[key_b] + alpha[key_a], flg_st_1);
				flg_st_1 = false;

				//alpha[s][t][1][1]
				key_b = getKey(r, s - 1, 0, 1, length);
				key_a = getKey(r, t, 0, 0, length);
				alpha[key_ts_1] = logsumexp(alpha[key_ts_1], beta[key_b] + alpha[key_a] + prodProb_rt, flg_ts_1);
				flg_ts_1 = false;

				key_a = getKey(r, t, 1, 0, length);
				alpha[key_ts_1] = logsumexp(alpha[key_ts_1], beta[key_b] + alpha[key_a] + prodProb_tr, flg_ts_1);
				flg_ts_1 = false;
			}

			for(int r = t + 1; r < length; r++){
				double prodProb_sr = probs[s][r][0];
				double prodProb_rs = probs[s][r][1];

				//alpha[s][t][0][1]
				int key_b = getKey(t + 1, r, 1, 1, length);
				int key_a = getKey(s, r, 0, 0, length);
				alpha[key_st_1] = logsumexp(alpha[key_st_1], beta[key_b] + alpha[key_a] + prodProb_sr, flg_st_1);
				flg_st_1 = false;

				key_a = getKey(s, r, 1, 0, length);
				alpha[key_st_1] = logsumexp(alpha[key_st_1], beta[key_b] + alpha[key_a] + prodProb_rs, flg_st_1);
				flg_st_1 = false;

				//alpha[s][t][1][1]
				key_b = getKey(t, r, 1, 0, length);
				key_a = getKey(s, r, 1, 1, length);
				alpha[key_ts_1] = logsumexp(alpha[key_ts_1], beta[key_b] + alpha[key_a], flg_ts_1);
				flg_ts_1 = false;
			}

			//alpha[s][t][0][0]
			for(int r = t; r < length; r++){
				int key_b = getKey(t, r, 0, 1, length);
				int key_a = getKey(s, r, 0, 1, length);
				alpha[key_st_0] = logsumexp(alpha[key_st_0], beta[key_b] + alpha[key_a], flg_st_0);
				flg_st_0 = false;
			}
			//alpha[s][t][1][0]
			for(int r = 0; r <= s; r++){
				int key_b = getKey(r, s, 1, 1, length);
				int key_a = getKey(r, t, 1, 1, length);
				alpha[key_ts_0] = logsumexp(alpha[key_ts_0], beta[key_b] + alpha[key_a], flg_ts_0);
				flg_ts_0 = false;
			}
		}
	}
}

void DependencyEncoder::updataExpected(long double *expected, int *keys, double m, int size){
	for(int i = 0; i < size; i++){
		int j = keys[i];
		expected[j] += m;
	}
}

void DependencyEncoder::getExpected(long double *expected, double *beta, double *alpha, double z, int length){
	// Get production crap.
	for (int w1 = 0; w1 < length; w1++) {
		for (int w2 = w1 + 1; w2 < length; w2++) {
			for (int ph = 0; ph < 2; ph++) {
				int *keys = NULL;
				int size = in->readArray(keys);
				int key = getKey(w1, w2, ph, 0, length);
				double m = exp(beta[key]  + alpha[key] - z);
				updataExpected(expected, keys, m, size);
				delete[](keys);
			}
		}
	}
	int last = in->readInt();
	if(last != -3){
		cout<<"Error reading file."<<endl;
		exit(0);
	}

	int *keys = NULL;
	int size = in->readArray(keys);
	for(int i = 0; i < size; i++){
		int j = keys[i];
		expected[j] -= 1.0;
	}
	
	delete[](keys);
	last = in->readInt();
	if (last != -4) {
		cout<<"Error reading file."<<endl;
		exit(0);
	}

	DependencyInstance* inst = new DependencyInstance();
	inst->readObject(*in);

	last = in->readInt();
	if (last != -1) {
		cout<<"Error reading file."<<endl;
		exit(0);
	}
	delete(inst);
}

double DependencyEncoder::gradient(long double *expected, int length){
	long offset = in->objTell();
	DependencyInstance *inst = pipe->readInstance(*in, length, probs, nt_probs, param);
	if(0 != (in->objSeek(offset))){
		cout<<"error on fseek of trainforset!"<<endl;
		exit(0);
	}

	double obj = 0.0;
	double *beta = new double[length * length * 2 * 2];
	double *alpha = new double[length * length * 2 * 2];

	double z = getZ(length, beta);

	inside_outside(length, beta, alpha);

	obj = z - param->getScore(inst->fv);
	delete(inst);

	getExpected(expected, beta, alpha, z, length);

	delete[](beta);
	delete[](alpha);

	return obj;
}

